/**
 * @author: mattwang@tencent.com
 * @date: 2012-10-16
 */

#include <ostream>
#include <iostream>
#include <fstream>
#include <math.h>

#include <map>
#include <string>
#include	 <type_def.h>
#include <param_def.h>
#include <maxent_core.h>

int MaxEntCore::train(ClassLoader &data) {
	LBFGSTrainer trainer;
	double sigma2 = 0.1;
	int max_iter = 100;
	double tol = 1e-5;
	prepare_data(data, trainer, sigma2);

	trainer.train(max_iter, tol);

	destroy_data(trainer);
}

int MaxEntCore::predict(ClassLoader &data, string infer_result) {
	return 0;
}

int MaxEntCore::prepare_data(ClassLoader &data, LBFGSTrainer& trainer, double sigma2) {
	trainer.train_data = &data.data_matrix;
	trainer.train_data_class = &data.data_class;
	trainer.train_data_weight = &data.data_weight;

	trainer.test_data = &data.test_matrix;
	trainer.test_data_class = &data.test_class;
	trainer.test_data_weight = &data.test_weight;

	trainer.n_train_data = trainer.train_data->size();
	trainer.n_test_data = trainer.test_data->size();

	int n_local_class = data.max_class_id;
	int n_global_class = 0;

	int n_feature = data.max_feature_id;
	int n_global_feature = 0;

	int info;
	MPI_Barrier (MPI_COMM_WORLD);
	info = MPI_Allreduce(&n_local_class, &n_global_class, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD);
	CHKERRQ(info);

	MPI_Barrier(MPI_COMM_WORLD);
	info = MPI_Allreduce(&n_feature, &n_global_feature, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD);
	CHKERRQ(info);

	trainer.n_class = n_global_class;
	Bitmap *class_hit = new Bitmap[n_global_class];
	for (int i = 0; i < n_global_class; i++) {
		class_hit[i].resize(n_global_feature);
	}

	for (int i = 0; i < trainer.n_train_data; i++) {
		SampleVector &sample = data.data_matrix.at(i);
		int class_id = data.data_class.at(i);

		for (size_t j = 0; j < sample.size(); ++j) {
			int fid = sample[i].id;
			class_hit[class_id].set(fid);
		}
	}

	Bitmap tmp_bitmap(n_global_feature);
	int n_theta = 0;
	for (int i = 0; i < n_global_class; i++) {
		int info;
		info = MPI_Allreduce(class_hit[i].buff(), tmp_bitmap.buff(), tmp_bitmap.real_size, MPI_BYTE, MPI_BXOR,
				MPI_COMM_WORLD);
		CHKERRQ(info);

		info = MPI_Barrier(MPI_COMM_WORLD);
		CHKERRQ(info);

		memcpy(tmp_bitmap.buff(), class_hit[i].buff(), tmp_bitmap.real_size);

		for (int j = 0; j < n_global_feature; j++) {
			if (tmp_bitmap.get(j) == 1) {
				n_theta++;
			}
		}
	}

	trainer.n_theta = n_theta;
	trainer.theta_all = new double[n_theta];

	std::fill(trainer.theta_all, trainer.theta_all + n_theta, 0.0);

	if (sigma2 > 0.0) {
		trainer.sigma2 = new double[n_theta];
	} else {
		trainer.sigma2 = NULL;
	}

	n_theta = 0;
	trainer.param_map = new vector<vector<pair<size_t, size_t> > >;
	trainer.param_map->reserve(n_global_feature);
	for (int fi = 0; fi < n_global_feature; fi++) {
		vector<pair<size_t, size_t> > &param = trainer.param_map->at(fi);

		for (int ci = 0; ci < n_global_class; ci++) {
			if (class_hit[ci].get(fi) == 1) {
				param.push_back(make_pair<size_t, size_t>(ci, n_theta++));
			}
		}
	}

	delete[] class_hit;

	return 0;
}
int MaxEntCore::destroy_data(LBFGSTrainer& trainer) {
	delete trainer.param_map;
	delete[] trainer.sigma2;
	delete[] trainer.theta_all;
	return 0;
}

